#include "cuDevice.h"
#include "cuImage.h"
#include "cuImage_filters.h"

#include "cuFunction.h"

#include "..\glew\glew.h"
#include <cuda_gl_interop.h>

#include <stdio.h>

#ifdef _WIN32
#   pragma warning( disable : 4996 ) // disable deprecated warning 
#endif  

//////////////////////////////////////////////////////////////////////////
// cuImage Ctor
//////////////////////////////////////////////////////////////////////////

cuImage_RGBA8::cuImage_RGBA8()
{
  Null();
}

//////////////////////////////////////////////////////////////////////////
// cuImage Dtor
//////////////////////////////////////////////////////////////////////////

cuImage_RGBA8::~cuImage_RGBA8()
{
  Release();
}

//////////////////////////////////////////////////////////////////////////
// cuImage Null function sets member variables to NULL or "0"
//////////////////////////////////////////////////////////////////////////

void cuImage_RGBA8::Null()
{
  m_pH = NULL;
  m_pD = NULL;
  m_pA = NULL;

  gl_pbo = gl_tex = 0;

  w = h = p = 0;
}

//////////////////////////////////////////////////////////////////////////
// cuImage setSize function sets size variables and computes pitch value
//////////////////////////////////////////////////////////////////////////

void cuImage_RGBA8::setSize(uint32 width, uint32 height)
{
  int mem_pitch = device_properties.textureAlignment;
  int elem_pitch = mem_pitch / sizeof(uchar4);

  // calculate pitch for the image
  this->w = width;
  this->h = height;
  this->p = width;
  
  if (width % elem_pitch) this->p += (elem_pitch - width % elem_pitch);
}

//////////////////////////////////////////////////////////////////////////
// cuImage Release functions
//////////////////////////////////////////////////////////////////////////

bool cuImage_RGBA8::Release(uint8 dFlags)
{
  cudaError_t cuda_error;

  if (m_pH != NULL) CUDA_CHECK_ERROR( cudaFreeH ( m_pH ) );
  if (m_pD != NULL) CUDA_CHECK_ERROR( cudaFreeD ( m_pD ) );
  if (m_pA != NULL) CUDA_CHECK_ERROR( cudaFreeA ( m_pA ) );

  if (gl_pbo) glDeleteBuffers (1, &gl_pbo);
  if (gl_tex) glDeleteTextures(1, &gl_tex);

  Null();

  if (glGetError() != GL_NO_ERROR)
  {
#ifdef _DEBUG
    fprintf(stderr, "Error while releasing PBO!\n");
#endif // _DEBUG
    return false;
  }

  return cuda_error == cudaSuccess;
}

bool cuImage_RGBA8::Release_Host()
{
  cudaError_t cuda_error;

  CUDA_CHECK_ERROR( cudaFreeH( m_pH) );

  m_pH = NULL;

  return cuda_error == cudaSuccess;
}

bool cuImage_RGBA8::Release_CUDA(uint8 cFlags /* = CU_IMAGE_D_MEMORY */)
{
  cudaError_t cuda_error;

  if (cFlags & CU_IMAGE_A_MEMORY)
  {
    CUDA_CHECK_ERROR( cudaFreeA( m_pA ) );

    m_pA = NULL;

    return cuda_error == cudaSuccess;
  }

  if (cFlags & CU_IMAGE_D_MEMORY)
  {
    CUDA_CHECK_ERROR( cudaFreeD( m_pD ) );

    m_pD = NULL;

    return cuda_error == cudaSuccess;
  }

  return false;
}

bool cuImage_RGBA8::Release_OGL()
{
  if (gl_pbo) glDeleteBuffers (1, &gl_pbo);
  if (gl_tex) glDeleteTextures(1, &gl_tex);

  gl_tex = gl_pbo = 0;
  
  if (glGetError() != GL_NO_ERROR)
  {
#ifdef _DEBUG
    fprintf(stderr, "Error while releasing PBO!\n");
#endif // _DEBUG
    return false;
  }

  return true;
}

//////////////////////////////////////////////////////////////////////////
// cuImage Create functions
//////////////////////////////////////////////////////////////////////////

bool cuImage_RGBA8::Create(uint32 width, 
                           uint32 height, 
                           uint8  cFlags)
{
  bool result = true;

  setSize(width, height);

  if ( cFlags & CU_IMAGE_H_MEMORY )  
       result = result && Create_Host();  

  if ( cFlags & CU_IMAGE_D_MEMORY ||
       cFlags & CU_IMAGE_A_MEMORY )
       result = result && Create_CUDA(cFlags, false);

  if ( cFlags & CU_IMAGE_OGL_RESOURCE )
       result = result && Create_OGL();

  return result;  
}

bool cuImage_RGBA8::Create_Host()
{
  cudaError_t cuda_error;

  // try allocating HOST resources 
  CUDA_CHECK_ERROR( cudaMallocH ( (void **) &m_pH, 
                                   p * h * sizeof(uchar4) ) );

  if ( m_pH == NULL )
  {
#ifdef _DEBUG    
    fprintf( stderr, "Error while allocating host memory \n" );
#endif // _DEBUG
    return false;
  }
  return cuda_error == cudaSuccess;
}

bool cuImage_RGBA8::Create_OGL()
{
  cudaError_t cuda_error;

  unsigned int pbo = 0;

  glGenBuffers(1, (GLuint *) &pbo);
  glBindBuffer(GL_PIXEL_UNPACK_BUFFER_ARB, pbo);
  glBufferData(GL_PIXEL_UNPACK_BUFFER_ARB, p * h * sizeof(uchar4), NULL, GL_STREAM_COPY);

  if (glGetError() != GL_NO_ERROR)
  {
#ifdef _DEBUG
    fprintf(stderr, "Error while creating PBO!\n");
#endif // _DEBUG
    return false;
  }

  CUDA_CHECK_ERROR( cudaGLRegisterBufferObject(pbo) );

  unsigned int tex = 0;
  fprintf(stderr, "Creating GL texture...\n");

  glEnable(GL_TEXTURE_2D);
  glGenTextures(1, (GLuint *) &tex);
  glBindTexture  (GL_TEXTURE_2D, tex);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S,   GL_CLAMP);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T,   GL_CLAMP);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  glTexImage2D   (GL_TEXTURE_2D, 0, GL_RGBA8, p, h, 0, GL_RGBA, GL_UNSIGNED_BYTE, NULL);

  if (glGetError() != GL_NO_ERROR) 
  {
#ifdef _DEBUG
    printf("Error while creating texture!\n");
#endif // _DEBUG
    return false;
  }

  gl_pbo = pbo;
  gl_tex = tex;

  return true;
}

bool cuImage_RGBA8::Create_CUDA(uint8 cFlags, 
                                bool  bCopy)
{
  cudaError_t cuda_error;

  if (cFlags & CU_IMAGE_A_MEMORY)
  {
    cudaChannelFormatDesc uchar4Desc = cudaCreateChannelDesc<uchar4>();

    CUDA_CHECK_ERROR( cudaMallocA( &m_pA, &uchar4Desc, p, h ) );

    if (m_pA == NULL || cuda_error != cudaSuccess )
      return false;

    if (bCopy && (m_pH != NULL))
      CUDA_CHECK_ERROR( cudaMemcpyHtoA(m_pA, m_pH, p * h * sizeof(uchar4)) );

    return cuda_error == cudaSuccess;
  }

  if (cFlags & CU_IMAGE_D_MEMORY)
  {
    CUDA_CHECK_ERROR( cudaMallocD( (void **) &m_pD, p * h * sizeof(uchar4) ) );

    if (m_pD == NULL || cuda_error != cudaSuccess)
      return false;

    if (bCopy && (m_pH != NULL))
      CUDA_CHECK_ERROR( cudaMemcpyHtoD( m_pD, m_pH, p * h * sizeof(uchar4)) );

    return cuda_error == cudaSuccess;
  }

  return false;
}


//////////////////////////////////////////////////////////////////////////
// cuImage "dynamic" filters: resource creation happens inside the call
//////////////////////////////////////////////////////////////////////////
cuImage_RGBA8 * cuImage_RGBA8::Clone()
{
  cuImage_RGBA8 * pDst = new cuImage_RGBA8;

  bool result = Clone(pDst, this);

#ifdef _DEBUG 
  if (!result) fprintf(stderr, "Error cloning images!\n");
#endif // _DEBUG

  return pDst;
}

cuImage_RGBA8 * cuImage_RGBA8::Gaussian(float sigma, 
                                        uint8 cFlags,
                                        uint8 rFlags,
                                        uint8 wFlags)
{
  cuImage_RGBA8 * pDst = new cuImage_RGBA8;

  pDst->Create(this->getW(), this->getH(), cFlags);

  bool result = Gaussian(pDst, this, sigma, 
                         rFlags, wFlags);

#ifdef _DEBUG
  if (!result) fprintf(stderr, "Error filtering (Gaussian) images!\n");
#endif // _DEBUG

  return pDst;
}

//////////////////////////////////////////////////////////////////////////
// cuImage "static" filters: resource creation happens before the call
//////////////////////////////////////////////////////////////////////////

bool cuImage_RGBA8::Clone(cuImage_RGBA8 * pDst, cuImage_RGBA8 * pSrc)
{
  if (pDst == NULL || pSrc == NULL || pDst == pSrc)
    return false;

  cudaError_t cuda_error;

  pDst->Release(); // release any previous content from pDst;

  uint8 createFlag = (pSrc->getHPtr() != NULL) * CU_IMAGE_H_MEMORY | 
                     (pSrc->getDPtr() != NULL) * CU_IMAGE_D_MEMORY |
                     (pSrc->getAPtr() != NULL) * CU_IMAGE_A_MEMORY |
                     (pSrc->getGlTex() && pSrc->getGlPBO()) * CU_IMAGE_OGL_RESOURCE;

  bool result = pDst->Create(pSrc->getW(), pSrc->getH(), createFlag);

  if (!result)
    return false;

  if (pSrc->getHPtr() != NULL)
    CUDA_CHECK_ERROR( cudaMemcpyHtoH( pDst->getHPtr(), pSrc->getHPtr(), 
                                      pDst->p * pDst->h * sizeof(uchar4)) );
 
  if (pDst->getDPtr() != NULL)
    CUDA_CHECK_ERROR( cudaMemcpyDtoD(pDst->getDPtr(), pSrc->getDPtr(),
                                     pDst->p * pDst->h * sizeof(uchar4)) );

  if (pDst->getAPtr())
    CUDA_CHECK_ERROR( cudaMemcpyAtoA( pDst->getAPtr(), pSrc->getAPtr(),
                                      pDst->p * pDst->h * sizeof(uchar4)) );

  return cuda_error == cudaSuccess;
}

bool cuImage_RGBA8::Gaussian(cuImage_RGBA8 * pDst, 
                             cuImage_RGBA8 * pSrc, 
                             float sigma,
                             uint8 rFlags,
                             uint8 wFlags)
{
  if (pDst == NULL || pSrc == NULL || pDst == pSrc)
    return false;

  cudaError_t cuda_error;

  //Release(); // uncertain if programmer should control this...

  return true;
}

bool cuImage_RGBA8::LoadBMP(const char * pFileName, 
                            uint8        createFlags)
{
  BMPHeader     hdr;
  BMPInfoHeader infoHdr;
  uint32        x, y;
  FILE         *fd = NULL;

#ifdef _DEBUG
  fprintf(stdout, "Loading %s...\n", pFileName);
#endif // _DEBUG

  //Release(); // uncertain if programmer should control this...

  if (sizeof(uchar4) != 4) {
#ifdef _DEBUG
    fprintf( stderr, "\t***Bad uchar4 size*** \n" );
#endif // _DEBUG
    return false;
  }

  if ( !( fd = fopen(pFileName, "rb") ) ) {
#ifdef _DEBUG
    fprintf( stderr, "\t***BMP load error: file %s access denied*** \n", pFileName );
#endif // _DEBUG
    return false;
  }

  fread( &hdr, sizeof(hdr), 1, fd );

  if ( hdr.type != 0x4D42 ) {
#ifdef _DEBUG
    fprintf( stderr, "\t***BMP load error: bad file format*** \n" );
#endif // _DEBUG
    return false;
  }

  fread( &infoHdr, sizeof(infoHdr), 1, fd );

  if ( infoHdr.bitsPerPixel != 24 ) {
#ifdef _DEBUG
    fprintf( stderr, "\t***BMP load error: invalid color depth*** \n" );
#endif // _DEBUG
    return false;
  }

  if ( infoHdr.compression ) {
#ifdef _DEBUG
    fprintf( stderr, "\t***BMP load error: compressed image*** \n" );
#endif // _DEBUG
    return false;
  }

  if ( !Create(infoHdr.width, infoHdr.height, createFlags) )
    return false;

#ifdef _DEBUG
  fprintf( stdout, "\tBMP width:  %u\n", w );
  fprintf( stdout, "\tBMP height: %u\n", h );
  fprintf( stdout, "\tBMP pitch:  %u\n", p );
#endif // _DEBUG

  fseek(fd, hdr.offset - sizeof(hdr) - sizeof(infoHdr), SEEK_CUR);

  for(y = 0; y < h; y++)
  {
    for(x = 0; x < p; x++)
    {
      if (x < w)
      {
        m_pH[(y * p + x)].w = 0xFF;      // set alpha channel to 255
        m_pH[(y * p + x)].z = fgetc(fd);
        m_pH[(y * p + x)].y = fgetc(fd);
        m_pH[(y * p + x)].x = fgetc(fd);
      }
      else
      {
        m_pH[(y * p + x)].w = 0xFF;      // set alpha channel to 255
        m_pH[(y * p + x)].z = 0;         // set rgb to 0 
        m_pH[(y * p + x)].y = 0;         // 
        m_pH[(y * p + x)].x = 0;         // 
      }
    }

    for(x = 0; x < (4 - (3 * w) % 4) % 4; x++) fgetc(fd); // read extra elements in case file is padded
  }


  if(ferror(fd)) {
#ifdef _DEBUG
    printf("\t***Unknown BMP load error.***\n");
#endif // _DEBUG
    Release();
    return false;
  }
#ifdef _DEBUG
  else printf("BMP file loaded successfully!\n");
#endif // _DEBUG

  fclose(fd);

  return true;
}

cudaDeviceProp cuImage_RGBA8::device_properties;
void cuImage_RGBA8::InitDeviceProperties()
{
  cuDevice(device_properties);
}
